# link: https://mp.weixin.qq.com/s/rFaU96l4mYzC0Kaua9jRJA

# 官方文档：https://www.backtrader.com/docu/indautoref/
'''
在编写策略时，除了常规的高开低收成交量等行情数据外，还会用到各式各样的指标（变量），
比如宏观经济指标、基本面分析指标、技术分析指标、另类数据等等。Backtrader 大致有 2 种获取指标的方式：

1、直接通过 DataFeeds 模块导入已经计算好的指标，比如《数据篇》中的导入新增指标 PE、PB；
2、在编写策略时调用 Indicators 指标模块临时计算指标，比如 5 日均线、布林带等 。
'''

# %%
import backtrader as bt
import backtrader.indicators as btind  # 导入策略分析模块
import pandas as pd
import datetime
import akshare as ak


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


codes = ['sh600276', 'sh600519', 'sh603288']

# 恒瑞医药
data1 = get_data(codes[0])
# 贵州茅台
data2 = get_data(codes[1])
# 海天味业
data3 = get_data(codes[2])

# 第1节 建议在 __init__() 中提前计算指标
'''
Strategy 中的 __init__() 函数在回测过程中只会在最开始的时候调用一次，而 next() 会每个交易日依次循环调用多次，
所以为了提高回测效率，建议先在 __init__() 中一次性计算好指标（甚至是交易信号），然后在 next() 中调用已经算好的指标，
这样能有效避免指标的重复计算，提高回测运行速度。

建议遵循“__init__() 负责指标计算，next() 负责指标调用”的原则。
'''


class TestStrategy(bt.Strategy):
    def __init__(self):
        # 指标函数也支持简写 SimpleMovingAverage → SMA
        # 最简方式：直接省略指向的数据集
        self.sma1 = btind.SimpleMovingAverage(period=5)
        # 只指定第一个数据表格
        self.sma2 = btind.SMA(self.data, period=5)
        # 指定第一个数据表格的 close 线
        self.sma3 = btind.SMA(self.data.close, period=5)
        # 完整写法
        self.sma4 = btind.SMA(self.datas[0].lines[0], period=5)

    def next(self):
        # 提取当前时间点
        print('datetime', self.datas[0].datetime.date(0))
        # 打印当日、昨日、前日的均线
        print('sma1', self.sma1.get(ago=0, size=3))
        print('sma2', self.sma2.get(ago=0, size=3))
        print('sma3', self.sma3.get(ago=0, size=3))
        print('sma4', self.sma4.get(ago=0, size=3))


cerebro = bt.Cerebro()
st_date = datetime.datetime(2023, 4, 1)
end_date = datetime.datetime(2023, 5, 12)
datafeed1 = bt.feeds.PandasData(dataname=data1, fromdate=st_date, todate=end_date)
cerebro.adddata(datafeed1, name=codes[0])
datafeed2 = bt.feeds.PandasData(dataname=data2, fromdate=st_date, todate=end_date)
cerebro.adddata(datafeed2, name=codes[1])

cerebro.addstrategy(TestStrategy)
result = cerebro.run()